import torch
from net import Model
from transformers import BertTokenizer

#定义设备信息
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)

pretrained = BertTokenizer.from_pretrained(
    r"D:\AIModels\llm\bert-base-chinese\models--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
    )
names = ["负向评价","正向评价"]
model = Model().to(DEVICE)

def collate_fn(data):
    """
    对数据进行编码
    :param data:
    :return:
    """
    sents = []
    sents.append(data)
    #编码
    data = pretrained.batch_encode_plus(
        batch_text_or_text_pairs=sents, #输入的文本
        truncation=True, #当句子长度大于max_length时，截断
        max_length=500, #设定最大序列长度
        padding="max_length", #填充到最大的序列长度， 不足的部分用0填充
        return_tensors="pt", #返回的类型, 如果为None，返回list， 如果为tf，返回tf张量， 如果为pt，返回pytorch张量， 如果为np，返回numpy数组
        return_length=True #返回length长度
    )
    input_ids = data["input_ids"]
    attention_mask = data["attention_mask"]
    token_type_ids = data["token_type_ids"]

    return input_ids, attention_mask, token_type_ids

def test():
    #加载训练参数
    model.load_state_dict(torch.load("params/best_bert.pth",map_location=DEVICE))
    #开启测试模式
    model.eval()

    while True:
        data = input("请输入测试数据（输入‘q’退出）：")
        if data == 'q':
            print("测试结束")
            break

        input_ids, attention_mask, token_type_ids = collate_fn(data)
        input_ids, attention_mask, token_type_ids = \
            input_ids.to(DEVICE), attention_mask.to(DEVICE), token_type_ids.to(DEVICE)

        with torch.no_grad():
            out = model(input_ids, attention_mask, token_type_ids)
            out = out.argmax(dim=1) #取最大值对应的索引, 0为负向评价，1为正向评价
            print("模型判定：",names[out],"\n")

if __name__ == '__main__':
    test()